{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp model_fn\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multitask Learning Model\n",
    "Multitask Learning Model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from typing import Dict, Tuple\n",
    "from inspect import signature\n",
    "from copy import copy\n",
    "\n",
    "from loguru import logger\n",
    "import tensorflow as tf\n",
    "from tensorflow.python.util.nest import flatten_with_joined_string_paths\n",
    "import transformers\n",
    "\n",
    "from m3tl.modeling import MultiModalBertModel\n",
    "from m3tl.params import Params\n",
    "from m3tl.utils import (get_embedding_table_from_model, get_transformer_main_model, create_dict_from_nested_model, variable_summaries, get_phase, set_phase)\n",
    "from m3tl.special_tokens import TRAIN, EVAL, PREDICT\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def filter_loss(loss, features, problem):\n",
    "\n",
    "    if tf.reduce_mean(input_tensor=features['%s_loss_multiplier' % problem]) == 0:\n",
    "        return_loss = 0.0\n",
    "    else:\n",
    "        return_loss = loss\n",
    "\n",
    "    return return_loss\n",
    "\n",
    "\n",
    "class BertMultiTaskBody(tf.keras.Model):\n",
    "    \"\"\"Model to extract bert features and dispatch corresponding rows to each problem_chunk.\n",
    "\n",
    "    for each problem chunk, we extract corresponding features\n",
    "    and hidden features for that problem. The reason behind this\n",
    "    is to save computation for downstream processing.\n",
    "    For example, we have a batch of two instances and they're from\n",
    "    problem a and b respectively:\n",
    "    Input:\n",
    "    [{'input_ids': [1,2,3], 'a_loss_multiplier': 1, 'b_loss_multiplier': 0},\n",
    "     {'input_ids': [4,5,6], 'a_loss_multiplier': 0, 'b_loss_multiplier': 1}]\n",
    "    Output:\n",
    "    {\n",
    "      'a': {'input_ids': [1,2,3], 'a_loss_multiplier': 1, 'b_loss_multiplier': 0}\n",
    "      'b': {'input_ids': [4,5,6], 'a_loss_multiplier': 0, 'b_loss_multiplier': 1}\n",
    "    }\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, params: Params, name='BertMultiTaskBody'):\n",
    "        super(BertMultiTaskBody, self).__init__(name=name)\n",
    "        self.params = params\n",
    "        self.bert = MultiModalBertModel(params=self.params)\n",
    "        if self.params.custom_pooled_hidden_size:\n",
    "            self.custom_pooled_layer = tf.keras.layers.Dense(\n",
    "                self.params.custom_pooled_hidden_size, activation=tf.keras.activations.selu)\n",
    "        else:\n",
    "            self.custom_pooled_layer = None\n",
    "\n",
    "    # @tf.function\n",
    "    def get_features_for_problem(self, features, hidden_feature, problem):\n",
    "        # get features with ind == 1\n",
    "        mode = get_phase()\n",
    "        if mode == PREDICT:\n",
    "            feature_this_round = features\n",
    "            hidden_feature_this_round = hidden_feature\n",
    "        else:\n",
    "            multiplier_name = '%s_loss_multiplier' % problem\n",
    "\n",
    "            # record_ind = tf.where(tf.cast(\n",
    "            #     tf.squeeze(features[multiplier_name]), tf.bool))\n",
    "\n",
    "            record_ind = tf.where(tf.cast(features[multiplier_name], tf.bool))\n",
    "\n",
    "            hidden_feature_this_round = {}\n",
    "            for hidden_feature_name in hidden_feature:\n",
    "                if hidden_feature_name != 'embed_table':\n",
    "                    hidden_feature_this_round[hidden_feature_name] = tf.squeeze(tf.gather(\n",
    "                        hidden_feature[hidden_feature_name], record_ind, axis=0\n",
    "                    ), axis=1)\n",
    "                    hidden_feature_this_round[hidden_feature_name].set_shape(\n",
    "                        [None, *hidden_feature[hidden_feature_name].shape.as_list()[1:]])\n",
    "                else:\n",
    "                    hidden_feature_this_round[hidden_feature_name] = hidden_feature[hidden_feature_name]\n",
    "\n",
    "            feature_this_round = {}\n",
    "            for features_name in features:\n",
    "                feature_this_round[features_name] = tf.gather_nd(\n",
    "                    features[features_name],\n",
    "                    record_ind)\n",
    "\n",
    "        return feature_this_round, hidden_feature_this_round\n",
    "\n",
    "    def call(self, inputs: Dict[str, tf.Tensor]) -> Tuple[Dict[str, Dict[str, tf.Tensor]], Dict[str, Dict[str, tf.Tensor]]]:\n",
    "        mode = get_phase()\n",
    "        inputs = copy(inputs)\n",
    "        # keys: ['last_hidden_state', 'pooler_output', 'hidden_states', \n",
    "        # 'attentions', 'model_input_mask', 'model_token_type_ids'\n",
    "        # model_input_mask, model_token_type_ids]\n",
    "        features, bert_model_outputs = self.bert(inputs, mode==TRAIN)\n",
    "\n",
    "        # extract bert hidden features\n",
    "        features['model_input_mask'] = bert_model_outputs['model_input_mask']\n",
    "        features['model_token_type_ids'] = bert_model_outputs['model_token_type_ids']\n",
    "\n",
    "        hidden_feature = {}\n",
    "        for logit_type in ['seq', 'pooled', 'all', 'embed', 'embed_table']:\n",
    "            if logit_type == 'seq':\n",
    "                # tensor, [batch_size, seq_length, hidden_size]\n",
    "                hidden_feature[logit_type] = bert_model_outputs['last_hidden_state']\n",
    "            elif logit_type == 'pooled':\n",
    "                # tensor, [batch_size, hidden_size]\n",
    "                hidden_feature[logit_type] = bert_model_outputs['pooler_output']\n",
    "                if self.custom_pooled_layer:\n",
    "                    hidden_feature[logit_type] = self.custom_pooled_layer(\n",
    "                        hidden_feature[logit_type])\n",
    "            elif logit_type == 'all':\n",
    "                # list, num_hidden_layers * [batch_size, seq_length, hidden_size]\n",
    "                hidden_feature[logit_type] = bert_model_outputs['all_encoder_layers']\n",
    "            elif logit_type == 'embed':\n",
    "                # for res connection\n",
    "                hidden_feature[logit_type] = bert_model_outputs['embedding_output']\n",
    "            elif logit_type == 'embed_table':\n",
    "                hidden_feature[logit_type] = bert_model_outputs['embedding_table']\n",
    "\n",
    "        if self.params.detail_log:\n",
    "            weight_var: tf.Variable\n",
    "            for weight_var in self.bert.weights:\n",
    "                variable_summaries(weight_var, 'body')\n",
    "\n",
    "        # for each problem chunk, we extract corresponding features\n",
    "        # and hidden features for that problem. The reason behind this\n",
    "        # is to save computation for downstream processing.\n",
    "        # For example, we have a batch of two instances and they're from\n",
    "        # problem a and b respectively:\n",
    "        # Input:\n",
    "        # [{'input_ids': [1,2,3], 'a_loss_multiplier': 1, 'b_loss_multiplier': 0},\n",
    "        #  {'input_ids': [4,5,6], 'a_loss_multiplier': 0, 'b_loss_multiplier': 1}]\n",
    "        # Output:\n",
    "        # {\n",
    "        #   'a': {'input_ids': [1,2,3], 'a_loss_multiplier': 1, 'b_loss_multiplier': 0}\n",
    "        #   'b': {'input_ids': [4,5,6], 'a_loss_multiplier': 0, 'b_loss_multiplier': 1}\n",
    "        # }\n",
    "        # NOTE: This logic now move to mtl model\n",
    "        return_feature = {}\n",
    "        return_hidden_feature = {}\n",
    "        return_feature['all'] = features\n",
    "        return_hidden_feature['all'] = hidden_feature\n",
    "        return return_feature, return_hidden_feature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-12 22:40:39.049 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_fake_ner, problem type: seq_tag\n",
      "2021-06-12 22:40:39.050 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-12 22:40:39.051 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_fake_cls, problem type: cls\n",
      "2021-06-12 22:40:39.051 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_masklm, problem type: masklm\n",
      "2021-06-12 22:40:39.051 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_fake_regression, problem type: regression\n",
      "2021-06-12 22:40:39.052 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_fake_vector_fit, problem type: vector_fit\n",
      "2021-06-12 22:40:39.052 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem weibo_premask_mlm, problem type: premask_mlm\n",
      "2021-06-12 22:40:39.053 | INFO     | m3tl.base_params:register_multiple_problems:526 - Adding new problem fake_contrastive_learning, problem type: contrastive_learning\n",
      "2021-06-12 22:40:39.053 | WARNING  | m3tl.base_params:assign_problem:620 - base_dir and dir_name arguments will be deprecated in the future. Please use model_dir instead.\n",
      "2021-06-12 22:40:39.054 | WARNING  | m3tl.base_params:prepare_dir:364 - bert_config not exists. will load model from huggingface checkpoint.\n",
      "2021-06-12 22:40:45.296 | WARNING  | m3tl.base_params:prepare_dir:364 - bert_config not exists. will load model from huggingface checkpoint.\n",
      "2021-06-12 22:40:45.366 | WARNING  | m3tl.read_write_tfrecord:chain_processed_data:248 - Chaining problems with & may consume a lot of memory if data is not pyspark RDD.\n",
      "2021-06-12 22:40:45.373 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpnowz79ap/weibo_fake_cls_weibo_fake_ner/train_00000.tfrecord\n",
      "2021-06-12 22:40:45.403 | WARNING  | m3tl.read_write_tfrecord:chain_processed_data:248 - Chaining problems with & may consume a lot of memory if data is not pyspark RDD.\n",
      "2021-06-12 22:40:45.409 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpnowz79ap/weibo_fake_cls_weibo_fake_ner/eval_00000.tfrecord\n",
      "2021-06-12 22:40:45.433 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpnowz79ap/weibo_fake_multi_cls/train_00000.tfrecord\n",
      "2021-06-12 22:40:45.458 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpnowz79ap/weibo_fake_multi_cls/eval_00000.tfrecord\n",
      "2021-06-12 22:40:45.535 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpnowz79ap/weibo_masklm/train_00000.tfrecord\n",
      "2021-06-12 22:40:45.585 | DEBUG    | m3tl.read_write_tfrecord:_write_fn:134 - Writing /tmp/tmpnowz79ap/weibo_masklm/eval_00000.tfrecord\n",
      "2021-06-12 22:40:46.375 | INFO     | m3tl.input_fn:train_eval_input_fn:56 - sampling weights: \n",
      "2021-06-12 22:40:46.376 | INFO     | m3tl.input_fn:train_eval_input_fn:57 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.35714285714285715,\n",
      "    \"weibo_fake_multi_cls\": 0.35714285714285715,\n",
      "    \"weibo_masklm\": 0.2857142857142857\n",
      "}\n",
      "2021-06-12 22:40:46.717 | INFO     | m3tl.input_fn:train_eval_input_fn:56 - sampling weights: \n",
      "2021-06-12 22:40:46.718 | INFO     | m3tl.input_fn:train_eval_input_fn:57 - {\n",
      "    \"weibo_fake_cls_weibo_fake_ner\": 0.35714285714285715,\n",
      "    \"weibo_fake_multi_cls\": 0.35714285714285715,\n",
      "    \"weibo_masklm\": 0.2857142857142857\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "from m3tl.test_base import TestBase\n",
    "import m3tl\n",
    "import shutil\n",
    "import numpy as np\n",
    "test_base = TestBase()\n",
    "test_base.params.assign_problem(\n",
    "    'weibo_fake_ner&weibo_fake_cls|weibo_fake_multi_cls|weibo_masklm')\n",
    "params = test_base.params\n",
    "train_dataset = m3tl.train_eval_input_fn(\n",
    "    params=params, mode=m3tl.TRAIN)\n",
    "eval_dataset = m3tl.train_eval_input_fn(\n",
    "    params=params, mode=m3tl.EVAL\n",
    ")\n",
    "\n",
    "train_dataset = train_dataset.repeat()\n",
    "\n",
    "one_batch_data = next(train_dataset.as_numpy_iterator())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "404 Client Error: Not Found for url: https://huggingface.co/voidful/albert_chinese_tiny/resolve/main/tf_model.h5\n",
      "Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFAlbertModel: ['predictions.LayerNorm.bias', 'predictions.LayerNorm.weight', 'predictions.dense.weight', 'predictions.decoder.weight', 'predictions.dense.bias', 'predictions.decoder.bias', 'predictions.bias']\n",
      "- This IS expected if you are initializing TFAlbertModel from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing TFAlbertModel from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of TFAlbertModel were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFAlbertModel for predictions without further training.\n",
      "2021-06-12 22:40:51.719 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: AutoGraph could not transform <bound method Socket.send of <zmq.sugar.socket.Socket object at 0x7f4a1c08b9f0>> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: module, class, method, function, traceback, frame, or code object was expected, got cython_function_or_method\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    }
   ],
   "source": [
    "mtl_body = BertMultiTaskBody(params=params)\n",
    "set_phase(TRAIN)\n",
    "features, hidden_features = mtl_body(one_batch_data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export \n",
    "class BertMultiTaskTop(tf.keras.Model):\n",
    "    \"\"\"Model to create top layer, aka classification layer, for each problem.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, params: Params, name='BertMultiTaskTop', input_embeddings: tf.Tensor = None):\n",
    "        super(BertMultiTaskTop, self).__init__(name=name)\n",
    "        self.params = params\n",
    "\n",
    "        problem_type_layer = self.params.top_layer\n",
    "        self.top_layer_dict = {}\n",
    "        for problem_dict in self.params.run_problem_list:\n",
    "            for problem in problem_dict:\n",
    "                problem_type = self.params.problem_type[problem]\n",
    "                # some layers has different signatures, assign inputs accordingly\n",
    "                layer_signature_name = signature(\n",
    "                    problem_type_layer[problem_type].__init__).parameters.keys()\n",
    "                inputs_kwargs = {\n",
    "                    'params': self.params,\n",
    "                    'problem_name': problem\n",
    "                }\n",
    "                for signature_name in layer_signature_name:\n",
    "                    if signature_name == 'input_embeddings':\n",
    "                        inputs_kwargs.update(\n",
    "                            {signature_name: input_embeddings})\n",
    "\n",
    "                self.top_layer_dict[problem] = problem_type_layer[problem_type](\n",
    "                    **inputs_kwargs)\n",
    "\n",
    "    def call(self,\n",
    "             inputs: Tuple[Dict[str, Dict[str, tf.Tensor]], Dict[str, Dict[str, tf.Tensor]]]) -> Dict[str, tf.Tensor]:\n",
    "        inputs = copy(inputs)\n",
    "        features, hidden_feature = inputs\n",
    "        return_dict = {}\n",
    "\n",
    "        mode = get_phase()\n",
    "\n",
    "        for problem_dict in self.params.run_problem_list:\n",
    "            for problem in problem_dict:\n",
    "                feature_this_round = features[problem]\n",
    "                hidden_feature_this_round = hidden_feature[problem]\n",
    "                problem_type = self.params.problem_type[problem]\n",
    "\n",
    "                # if pretrain, return pretrain logit\n",
    "                if problem_type == 'pretrain':\n",
    "                    pretrain = self.top_layer_dict[problem]\n",
    "                    return_dict[problem] = pretrain(\n",
    "                        (feature_this_round, hidden_feature_this_round))\n",
    "                    return return_dict\n",
    "\n",
    "                with tf.name_scope(problem):\n",
    "                    layer = self.top_layer_dict[problem]\n",
    "\n",
    "                    return_dict[problem] = layer(\n",
    "                        (feature_this_round, hidden_feature_this_round))\n",
    "\n",
    "        return return_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-12 22:40:58.059 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n"
     ]
    }
   ],
   "source": [
    "# use embedding table as the classification top of mask lm\n",
    "# top layer takes per-problem features and hidden features\n",
    "from m3tl.utils import dispatch_features\n",
    "\n",
    "features_per_problem, hidden_features_per_problem = {}, {}\n",
    "for problem in params.problem_list:\n",
    "    features_per_problem[problem], hidden_features_per_problem[problem] = dispatch_features(\n",
    "        features=features['all'], hidden_feature=hidden_features['all'], problem=problem, mode=TRAIN\n",
    "    )\n",
    "\n",
    "input_embeddings = m3tl.utils.get_embedding_table_from_model(\n",
    "    mtl_body.bert.bert_model)\n",
    "mtl_top = BertMultiTaskTop(params=params, input_embeddings=input_embeddings)\n",
    "set_phase(TRAIN)\n",
    "logit_dict = mtl_top((features_per_problem, hidden_features_per_problem))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for problem, problem_logit in logit_dict.items():\n",
    "    # last dim of logits equals to num_classes\n",
    "    assert problem_logit.shape[-1] == params.get_problem_info(problem=problem, info_name='num_classes')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export \n",
    "\n",
    "class BertMultiTask(tf.keras.Model):\n",
    "    def __init__(self, params: Params, name='BertMultiTask') -> None:\n",
    "        super(BertMultiTask, self).__init__(name=name)\n",
    "        self.params = params\n",
    "        # initialize body model, aka transformers\n",
    "        self.body = BertMultiTaskBody(params=self.params)\n",
    "\n",
    "        self.mtl_model = self.params.mtl_model['model'](\n",
    "            self.params, self.params.mtl_model['model_name'])\n",
    "        self.mtl_model_include_top = self.params.mtl_model['include_top']\n",
    "\n",
    "        # mlm might need word embedding from bert, map_structure\n",
    "        # build sub-model\n",
    "        input_embeddings = get_embedding_table_from_model(self.body.bert.bert_model)\n",
    "        self.top = BertMultiTaskTop(\n",
    "            params=self.params, input_embeddings=input_embeddings)\n",
    "\n",
    "        # loss combination strategy\n",
    "        self.loss_combination = self.params.loss_combination_strategy['model'](\n",
    "            self.params, self.params.loss_combination_strategy['name'])\n",
    "        self.eval_loss_metric_dict_list = []\n",
    "        self.eval_metric_dict_list = []\n",
    "        \n",
    "        # for horovod\n",
    "        self.first_batch = True\n",
    "\n",
    "    def call(self, inputs):\n",
    "        mode = get_phase()\n",
    "        inputs = copy(inputs)\n",
    "        pred_dict = {}\n",
    "        body_outputs = self.body(\n",
    "            inputs)\n",
    "        if self.params.output_body_seq_hidden:\n",
    "            pred_dict['seq'] = body_outputs[1]['all']['seq']\n",
    "            \n",
    "        if self.params.output_body_pooled_hidden:\n",
    "            pred_dict['pooled'] = body_outputs[1]['all']['pooled']\n",
    "            \n",
    "        if self.mtl_model_include_top:\n",
    "            problem_pred = self.mtl_model(\n",
    "                body_outputs)\n",
    "            pred_dict.update(problem_pred)\n",
    "            return pred_dict\n",
    "        \n",
    "        mtl_model_outputs = self.mtl_model(\n",
    "            body_outputs)\n",
    "        \n",
    "        if self.params.output_mtl_model_hidden:\n",
    "            pred_dict['mtl'] = mtl_model_outputs\n",
    "        \n",
    "        pred_per_problem = self.top(\n",
    "            (mtl_model_outputs))\n",
    "        \n",
    "        if self.params.output_problem_pred:\n",
    "            pred_dict.update(pred_per_problem)\n",
    "        return pred_dict\n",
    "\n",
    "    def compile(self, *args, **kwargs):\n",
    "        super(BertMultiTask, self).compile(*args, **kwargs)\n",
    "        logger.critical('Initial lr: {}'.format(self.params.lr))\n",
    "        logger.critical('Train steps: {}'.format(self.params.train_steps))\n",
    "        logger.critical('Warmup steps: {}'.format(self.params.num_warmup_steps))\n",
    "        self.optimizer, self.lr_scheduler = transformers.optimization_tf.create_optimizer(\n",
    "            init_lr=self.params.lr,\n",
    "            num_train_steps=self.params.train_steps,\n",
    "            num_warmup_steps=self.params.num_warmup_steps,\n",
    "            weight_decay_rate=0.01\n",
    "        )\n",
    "        self.mean_acc = tf.keras.metrics.Mean(name='mean_acc')\n",
    "\n",
    "    def add_flatten_losses_metrics(self, return_dict: dict):\n",
    "        current_eval_loss_dict = create_dict_from_nested_model(self, ele_name='losses')\n",
    "        flatten_losses = flatten_with_joined_string_paths(current_eval_loss_dict)\n",
    "        flatten_losses = {p: v for p, v in flatten_losses}\n",
    "        return_dict.update(flatten_losses)\n",
    "        return return_dict\n",
    "\n",
    "    def train_step(self, data):\n",
    "        set_phase(TRAIN)\n",
    "        with tf.GradientTape() as tape:\n",
    "            # Forward pass\n",
    "            _ = self(data)\n",
    "            # gather losses from all problems\n",
    "            metric_dict = {m.name: m.result() for m in self.metrics}\n",
    "            metric_dict = self.add_flatten_losses_metrics(return_dict=metric_dict)\n",
    "\n",
    "            # apply loss combination strategy\n",
    "            current_round_nest_loss_metric = create_dict_from_nested_model(self, ele_name='losses')\n",
    "            combined_losses = self.loss_combination(\n",
    "                current_loss_dict=current_round_nest_loss_metric,\n",
    "                current_metric_dict=self.metrics,\n",
    "                history=self.history\n",
    "            )\n",
    "            \n",
    "        if self.params.use_horovod:\n",
    "            import horovod.tensorflow as hvd\n",
    "            \n",
    "            tape = hvd.DistributedGradientTape(tape)\n",
    "        # Compute gradients\n",
    "        trainable_vars = self.trainable_variables\n",
    "        gradients = tape.gradient(combined_losses, trainable_vars)\n",
    "\n",
    "        # Update weights\n",
    "        self.optimizer.apply_gradients(zip(gradients, trainable_vars))\n",
    "\n",
    "        self.mean_acc.update_state(\n",
    "            [v for n, v in metric_dict.items() if n != 'mean_acc'])\n",
    "\n",
    "        return_dict = metric_dict\n",
    "        return_dict[self.mean_acc.name] = self.mean_acc.result()\n",
    "        # Return a dict mapping metric names to current value.\n",
    "        # Note that it will include the loss (tracked in self.metrics).\n",
    "        \n",
    "        if self.first_batch and self.params.use_horovod:\n",
    "            self.first_batch = False\n",
    "            hvd.broadcast_variables(trainable_vars, root_rank=0)\n",
    "            hvd.broadcast_variables(self.optimizer.variables(), root_rank=0)\n",
    "            \n",
    "        return return_dict\n",
    "\n",
    "    def test_step(self, data):\n",
    "        \"\"\"The logic for one evaluation step.\n",
    "\n",
    "        This method can be overridden to support custom evaluation logic.\n",
    "        This method is called by `Model.make_test_function`.\n",
    "\n",
    "        This function should contain the mathemetical logic for one step of\n",
    "        evaluation.\n",
    "        This typically includes the forward pass, loss calculation, and metrics\n",
    "        updates.\n",
    "\n",
    "        Configuration details for *how* this logic is run (e.g. `tf.function` and\n",
    "        `tf.distribute.Strategy` settings), should be left to\n",
    "        `Model.make_test_function`, which can also be overridden.\n",
    "\n",
    "        Arguments:\n",
    "        data: A nested structure of `Tensor`s.\n",
    "\n",
    "        Returns:\n",
    "        A `dict` containing values that will be passed to\n",
    "        `tf.keras.callbacks.CallbackList.on_train_batch_end`. Typically, the\n",
    "        values of the `Model`'s metrics are returned.\n",
    "        \"\"\"\n",
    "        set_phase(EVAL)\n",
    "        y_pred = self(data)\n",
    "        # Updates stateful loss metrics.\n",
    "        self.compiled_loss(\n",
    "            None, y_pred, None, regularization_losses=self.losses)\n",
    "\n",
    "        self.compiled_metrics.update_state(None, y_pred, None)\n",
    "\n",
    "        # get metrics to calculate mean\n",
    "        m_list = []\n",
    "        for metric in self.metrics:\n",
    "            if 'mean_acc' in metric.name:\n",
    "                continue\n",
    "            if 'acc' in metric.name:\n",
    "                m_list.append(metric.result())\n",
    "\n",
    "            if 'f1' in metric.name:\n",
    "                m_list.append(metric.result())\n",
    "\n",
    "        self.mean_acc.update_state(\n",
    "            m_list)\n",
    "\n",
    "        ret_dict = {m.name: m.result() for m in self.metrics}\n",
    "        ret_dict = self.add_flatten_losses_metrics(ret_dict)\n",
    "        return ret_dict\n",
    "        \n",
    "    def predict_step(self, data):\n",
    "        set_phase(PREDICT)\n",
    "        return self(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "404 Client Error: Not Found for url: https://huggingface.co/voidful/albert_chinese_tiny/resolve/main/tf_model.h5\n",
      "Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFAlbertModel: ['predictions.LayerNorm.bias', 'predictions.LayerNorm.weight', 'predictions.dense.weight', 'predictions.decoder.weight', 'predictions.dense.bias', 'predictions.decoder.bias', 'predictions.bias']\n",
      "- This IS expected if you are initializing TFAlbertModel from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing TFAlbertModel from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of TFAlbertModel were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFAlbertModel for predictions without further training.\n",
      "2021-06-12 22:41:03.366 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-12 22:41:03.459 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    }
   ],
   "source": [
    "mtl = BertMultiTask(params=params)\n",
    "logit_dict = mtl(one_batch_data)\n",
    "for problem, problem_logit in logit_dict.items():\n",
    "    # last dim of logits equals to num_classes\n",
    "    assert problem_logit.shape[-1] == params.get_problem_info(problem=problem, info_name='num_classes')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-12 22:41:04.946 | CRITICAL | __main__:compile:62 - Initial lr: 0.0\n",
      "2021-06-12 22:41:04.947 | CRITICAL | __main__:compile:63 - Train steps: 0\n",
      "2021-06-12 22:41:04.947 | CRITICAL | __main__:compile:64 - Warmup steps: 0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/3\n",
      "WARNING: AutoGraph could not transform <bound method BertMultiTaskBody.call of <__main__.BertMultiTaskBody object at 0x7f490c151cd0>> and will run it as-is.\n",
      "Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.\n",
      "Cause: invalid value for \"node\": expected \"ast.AST\", got \"<class 'NoneType'>\"; to visit lists of nodes, use \"visit_block\" instead\n",
      "To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/1 [==============================] - ETA: 0s - mean_acc: 2.4794 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.1143 - BertMultiTaskTop/weibo_fake_cls/losses/0: 0.9767 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.2088 - BertMultiTaskTop/weibo_fake_ner/losses/0: 2.0114 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0653"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/1 [==============================] - 14s 14s/step - mean_acc: 2.4794 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.1143 - BertMultiTaskTop/weibo_fake_cls/losses/0: 0.9767 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.2088 - BertMultiTaskTop/weibo_fake_ner/losses/0: 2.0114 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0653 - val_loss: 14.3237 - val_mean_acc: 0.3214 - val_weibo_fake_cls_acc: 0.5000 - val_weibo_fake_ner_acc: 0.1429 - val_BertMultiTaskTop/weibo_fake_cls/losses/0: 0.9819 - val_BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.2721 - val_BertMultiTaskTop/weibo_fake_ner/losses/0: 2.0120 - val_BertMultiTaskTop/weibo_masklm/losses/0: 10.0577\n",
      "Epoch 2/3\n",
      "1/1 [==============================] - 1s 717ms/step - mean_acc: 2.4846 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.1143 - BertMultiTaskTop/weibo_fake_cls/losses/0: 0.9992 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.1767 - BertMultiTaskTop/weibo_fake_ner/losses/0: 2.0816 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0360 - val_loss: 14.3237 - val_mean_acc: 0.3214 - val_weibo_fake_cls_acc: 0.5000 - val_weibo_fake_ner_acc: 0.1429 - val_BertMultiTaskTop/weibo_fake_cls/losses/0: 0.9819 - val_BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.2721 - val_BertMultiTaskTop/weibo_fake_ner/losses/0: 2.0120 - val_BertMultiTaskTop/weibo_masklm/losses/0: 10.0577\n",
      "Epoch 3/3\n",
      "1/1 [==============================] - 1s 602ms/step - mean_acc: 2.5813 - weibo_fake_cls_acc: 0.5000 - weibo_fake_ner_acc: 0.1429 - BertMultiTaskTop/weibo_fake_cls/losses/0: 1.5078 - BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.2698 - BertMultiTaskTop/weibo_fake_ner/losses/0: 2.0046 - BertMultiTaskTop/weibo_masklm/losses/0: 10.0627 - val_loss: 14.3237 - val_mean_acc: 0.3214 - val_weibo_fake_cls_acc: 0.5000 - val_weibo_fake_ner_acc: 0.1429 - val_BertMultiTaskTop/weibo_fake_cls/losses/0: 0.9819 - val_BertMultiTaskTop/weibo_fake_multi_cls/losses/0: 1.2721 - val_BertMultiTaskTop/weibo_fake_ner/losses/0: 2.0120 - val_BertMultiTaskTop/weibo_masklm/losses/0: 10.0577\n"
     ]
    }
   ],
   "source": [
    "mtl.compile()\n",
    "hist = mtl.fit(train_dataset, validation_data=train_dataset, validation_steps=1, steps_per_epoch=1, epochs=3, verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "404 Client Error: Not Found for url: https://huggingface.co/voidful/albert_chinese_tiny/resolve/main/tf_model.h5\n",
      "Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFAlbertModel: ['predictions.LayerNorm.bias', 'predictions.LayerNorm.weight', 'predictions.dense.weight', 'predictions.decoder.weight', 'predictions.dense.bias', 'predictions.decoder.bias', 'predictions.bias']\n",
      "- This IS expected if you are initializing TFAlbertModel from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing TFAlbertModel from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of TFAlbertModel were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFAlbertModel for predictions without further training.\n",
      "2021-06-12 22:41:25.001 | CRITICAL | m3tl.embedding_layer.base:__init__:58 - Modal Type id mapping: \n",
      " {\n",
      "    \"array\": 0,\n",
      "    \"cate\": 1,\n",
      "    \"text\": 2\n",
      "}\n",
      "2021-06-12 22:41:25.094 | WARNING  | m3tl.problem_types.masklm:__init__:41 - Share embedding is enabled but hidden_size != embedding_size\n",
      "The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).\n",
      "The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.\n"
     ]
    }
   ],
   "source": [
    "params.output_body_pooled_hidden = True\n",
    "params.output_body_seq_hidden = True\n",
    "params.output_mtl_model_hidden = True\n",
    "mtl = BertMultiTask(params=params)\n",
    "logit_dict = mtl(one_batch_data)\n",
    "assert 'pooled' in logit_dict\n",
    "assert 'seq' in logit_dict\n",
    "assert 'mtl' in logit_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.3 64-bit ('base': conda)",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
